import cv2
import os
import time
import random
import numpy as np
import pyautogui
import shutil
from ultralytics import YOLO
import yaml
import threading
from datetime import datetime

# 配置中文字体支持
pyautogui.FAILSAFE = False  # 关闭安全模式，避免鼠标移到角落停止程序


class DNFFiendishQiDataset:
    """DNF妖气追踪数据集处理类"""

    def __init__(self, data_root="dnf_yaqi_dataset"):
        self.data_root = data_root
        self.screenshot_dir = os.path.join(data_root, "raw_screenshots")
        self.train_img_dir = os.path.join(data_root, "images", "train")
        self.val_img_dir = os.path.join(data_root, "images", "val")
        self.train_label_dir = os.path.join(data_root, "labels", "train")
        self.val_label_dir = os.path.join(data_root, "labels", "val")

        # 创建必要的目录
        for dir_path in [self.screenshot_dir, self.train_img_dir, self.val_img_dir,
                         self.train_label_dir, self.val_label_dir]:
            os.makedirs(dir_path, exist_ok=True)

        # 妖气追踪相关的目标类别
        self.class_names = [
            "yaqi_mark",  # 妖气标记(地图上的红色/金色标记)
            "elite_monster",  # 精英怪(带有妖气的特殊怪物)
            "boss",  # BOSS(最终目标)
            "player",  # 玩家自身
            "portal",  # 传送门
            "item_yaqi",  # 妖气相关物品
            "mini_map",  # 小地图(用于定位妖气位置)
            "direction_arrow"  # 方向指引箭头
        ]

    def capture_in_game_screenshots(self, duration=300, interval=0.3):
        """捕获游戏内截图，专注于妖气相关场景"""
        print(f"开始捕获妖气场景截图，持续{duration}秒...")
        start_time = time.time()
    
        # 使用窗口标题获取位置（替代硬编码区域）
        try:
            # 使用pygetwindow库通过窗口标题查找游戏窗口
            import pygetwindow as gw
            
            # 查找包含"地下城与勇士"或"DNF"的窗口
            game_windows = [w for w in gw.getAllWindows() if 
                           ("地下城与勇士" in w.title or "DNF" in w.title) and
                           w.visible]
            
            if game_windows:
                # 选择第一个找到的游戏窗口
                game_window = game_windows[0]
                # 确保窗口在前台
                game_window.activate()
                time.sleep(0.5)  # 等待窗口切换
                
                # 获取窗口位置和大小
                x, y = game_window.left, game_window.top
                width, height = game_window.width, game_window.height
                
                # 设置截图区域
                game_region = (x, y, width, height)
                print(f"找到游戏窗口: {game_window.title}, 区域: {game_region}")
            else:
                # 如果没找到，回退到硬编码区域
                print("未找到游戏窗口，使用默认区域")
                game_region = (100, 100, 1024, 768)  # 回退到默认值
        except Exception as e:
            print(f"获取游戏窗口位置时出错: {str(e)}，使用默认区域")
            game_region = (100, 100, 1024, 768)  # 回退到默认值
    
        while time.time() - start_time < duration:
            # 捕获游戏窗口区域
            screenshot = pyautogui.screenshot(region=game_region)
            frame = cv2.cvtColor(np.array(screenshot), cv2.COLOR_RGB2BGR)

            # 保存截图
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
            save_path = os.path.join(self.screenshot_dir, f"yaqi_{timestamp}.jpg")
            cv2.imwrite(save_path, frame)

            # 每10秒打印一次进度
            if int(time.time() - start_time) % 10 == 0:
                print(f"已捕获 {int(time.time() - start_time)} 秒, 保存路径: {save_path}")

            time.sleep(interval)

        total = len(os.listdir(self.screenshot_dir))
        print(f"截图完成，共捕获 {total} 张妖气场景图片")

    def prepare_dataset(self, train_ratio=0.8):
        """准备训练集和验证集"""
        # 获取所有截图
        image_extensions = ('.jpg', '.jpeg', '.png')
        images = [f for f in os.listdir(self.screenshot_dir)
                  if f.lower().endswith(image_extensions)]

        if not images:
            print("未找到截图，请先执行截图捕获")
            return False

        # 随机打乱并分割
        random.shuffle(images)
        split_idx = int(len(images) * train_ratio)
        train_images = images[:split_idx]
        val_images = images[split_idx:]

        # 复制训练集
        for img in train_images:
            src_img = os.path.join(self.screenshot_dir, img)
            dst_img = os.path.join(self.train_img_dir, img)
            shutil.copy2(src_img, dst_img)

            # 复制对应的标注文件
            label_name = os.path.splitext(img)[0] + ".txt"
            src_label = os.path.join(self.screenshot_dir, label_name)
            if os.path.exists(src_label):
                dst_label = os.path.join(self.train_label_dir, label_name)
                shutil.copy2(src_label, dst_label)

        # 复制验证集
        for img in val_images:
            src_img = os.path.join(self.screenshot_dir, img)
            dst_img = os.path.join(self.val_img_dir, img)
            shutil.copy2(src_img, dst_img)

            # 复制对应的标注文件
            label_name = os.path.splitext(img)[0] + ".txt"
            src_label = os.path.join(self.screenshot_dir, label_name)
            if os.path.exists(src_label):
                dst_label = os.path.join(self.val_label_dir, label_name)
                shutil.copy2(src_label, dst_label)

        print(f"数据集准备完成 - 训练集: {len(train_images)} 张, 验证集: {len(val_images)} 张")
        return True

    def create_config_file(self):
        """创建YOLO训练所需的配置文件"""
        config = {
            "train": os.path.abspath(self.train_img_dir),
            "val": os.path.abspath(self.val_img_dir),
            "nc": len(self.class_names),
            "names": self.class_names
        }

        config_path = os.path.join(self.data_root, "yaqi_data.yaml")
        with open(config_path, "w", encoding="utf-8") as f:
            yaml.dump(config, f, sort_keys=False, allow_unicode=True)

        print(f"配置文件已创建: {config_path}")
        return config_path


class DNFFiendishQiTrainer:
    """DNF妖气追踪模型训练器"""

    def __init__(self, config_path):
        self.config_path = config_path
        # 选择适合游戏场景检测的模型
        self.model = YOLO("yolov8m.pt")  # 使用中等大小的模型平衡速度和精度

    def train(self, epochs=100, imgsz=640):
        """训练妖气追踪模型"""
        print("开始训练DNF妖气追踪模型...")

        # 检查GPU可用性
        device = 0 if cv2.cuda.getCudaEnabledDeviceCount() > 0 else 'cpu'
        print(f"训练设备: {'GPU' if device == 0 else 'CPU'}")

        # 开始训练，针对游戏场景优化参数
        results = self.model.train(
            data=self.config_path,
            epochs=epochs,
            imgsz=imgsz,
            device=device,
            batch=-1,  # 自动确定批次大小
            project="dnf_yaqi_models",
            name="yaqi_tracker",
            pretrained=True,
            augment=True,  # 启用数据增强提高泛化能力
            patience=20,  # 早停机制，防止过拟合
            lr0=0.001,  # 初始学习率
            weight_decay=0.0005  # 权重衰减，防止过拟合
        )

        print("妖气追踪模型训练完成!")
        return results


class DNFFiendishQiController:
    """DNF妖气追踪自动刷图控制器"""

    def __init__(self, model_path):
        self.model = YOLO(model_path)
        self.running = False
        self.class_names = self.model.names
        self.screen_center = (512, 384)  # 游戏窗口中心坐标(基于1024x768)

        # 技能和移动按键映射(可根据实际按键设置修改)
        self.keys = {
            "move_up": "up",
            "move_down": "down",
            "move_left": "left",
            "move_right": "right",
            "basic_attack": "j",
            "skill1": "a",  # 主要输出技能
            "skill2": "s",  # 控制技能
            "skill3": "d",  # 范围技能
            "ultimate": "f",  # 大招
            "pick_item": "space",  # 拾取物品
            "use_potion": "1",  # 使用药水
            "character_select": "esc"  # 角色选择快捷键
        }

        # 状态变量
        self.cooldowns = {
            "skill1": 0,
            "skill2": 0,
            "skill3": 0,
            "ultimate": 0,
            "potion": 0
        }

        # 冷却时间(秒)
        self.cooldown_times = {
            "skill1": 2.0,
            "skill2": 5.0,
            "skill3": 8.0,
            "ultimate": 30.0,
            "potion": 10.0
        }
        
        # 角色切换相关配置
        self.current_character_index = 0
        self.characters = ["character1", "character2", "character3"]  # 角色列表，可根据实际角色数量修改
        self.fatigue_threshold = 10  # 疲劳值阈值，低于此值时切换角色

    def get_game_frame(self):
        """获取游戏画面帧"""
        try:
            import pygetwindow as gw
            
            # 查找游戏窗口
            game_windows = [w for w in gw.getAllWindows() if 
                           ("地下城与勇士" in w.title or "DNF" in w.title) and 
                           w.visible]
            
            if game_windows:
                game_window = game_windows[0]
                x, y = game_window.left, game_window.top
                width, height = game_window.width, game_window.height
                # 更新屏幕中心坐标
                self.screen_center = (width // 2, height // 2)
                return pyautogui.screenshot(region=(x, y, width, height))
            else:
                print("未找到游戏窗口")
                return None
        except Exception as e:
            print(f"获取游戏画面时出错: {str(e)}")
            return None

    def detect_yaqi_targets(self, frame):
        """检测妖气相关目标"""
        # 较高的置信度阈值，确保检测准确性
        results = self.model(frame, conf=0.65)
        return results

    def is_on_cooldown(self, skill_name):
        """检查技能是否在冷却中"""
        return time.time() < self.cooldowns[skill_name]

    def use_skill(self, skill_name):
        """使用指定技能"""
        if self.is_on_cooldown(skill_name):
            return False

        # 按下技能键
        pyautogui.press(self.keys[skill_name])
        # 更新冷却时间
        self.cooldowns[skill_name] = time.time() + self.cooldown_times[skill_name]
        print(f"使用技能: {skill_name}")
        return True

    def move_towards(self, target_pos, threshold=50):
        """向目标位置移动"""
        tx, ty = target_pos
        cx, cy = self.screen_center

        # 水平方向移动
        if tx < cx - threshold:  # 目标在左侧
            pyautogui.keyDown(self.keys["move_left"])
            time.sleep(0.1)
            pyautogui.keyUp(self.keys["move_left"])
        elif tx > cx + threshold:  # 目标在右侧
            pyautogui.keyDown(self.keys["move_right"])
            time.sleep(0.1)
            pyautogui.keyUp(self.keys["move_right"])

        # 垂直方向移动
        if ty < cy - threshold:  # 目标在上侧
            pyautogui.keyDown(self.keys["move_up"])
            time.sleep(0.1)
            pyautogui.keyUp(self.keys["move_up"])
        elif ty > cy + threshold:  # 目标在下侧
            pyautogui.keyDown(self.keys["move_down"])
            time.sleep(0.1)
            pyautogui.keyUp(self.keys["move_down"])

    def yaqi_tracking_logic(self, frame, results):
        """妖气追踪核心逻辑"""
        # 解析检测结果
        yaqi_marks = []
        elite_monsters = []
        boss = None
        player = None
        portal = None

        for result in results:
            for box in result.boxes:
                cls_id = int(box.cls[0])
                cls_name = self.class_names[cls_id]
                conf = float(box.conf[0])

                # 计算目标中心位置
                x1, y1, x2, y2 = map(int, box.xyxy[0])
                center_x, center_y = (x1 + x2) // 2, (y1 + y2) // 2

                # 绘制检测框和标签
                color_map = {
                    "yaqi_mark": (0, 0, 255),  # 红色
                    "elite_monster": (0, 165, 255),  # 橙色
                    "boss": (128, 0, 128),  # 紫色
                    "player": (0, 255, 0),  # 绿色
                    "portal": (255, 0, 255)  # 粉色
                }
                color = color_map.get(cls_name, (255, 255, 0))  # 默认黄色

                cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
                cv2.putText(frame, f"{cls_name} {conf:.2f}",
                            (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX,
                            0.5, color, 2)

                # 分类目标
                if cls_name == "yaqi_mark":
                    yaqi_marks.append((center_x, center_y, conf))
                elif cls_name == "elite_monster":
                    elite_monsters.append((center_x, center_y, conf))
                elif cls_name == "boss":
                    boss = (center_x, center_y, conf)
                elif cls_name == "player":
                    player = (center_x, center_y)
                elif cls_name == "portal":
                    portal = (center_x, center_y)

        # 决策逻辑: 优先追踪妖气标记和精英怪
        action_taken = False

        # 1. 优先攻击BOSS
        if boss and player:
            self.move_towards((boss[0], boss[1]))
            # 使用大招攻击BOSS
            self.use_skill("ultimate")
            # 配合其他技能
            self.use_skill("skill3")
            self.use_skill("skill1")
            action_taken = True

        # 2. 攻击精英怪
        elif elite_monsters and player:
            # 选择置信度最高的精英怪
            elite_monsters.sort(key=lambda x: x[2], reverse=True)
            target = elite_monsters[0]
            self.move_towards((target[0], target[1]))
            # 使用技能攻击
            self.use_skill("skill2")
            self.use_skill("skill1")
            action_taken = True

        # 3. 追踪妖气标记
        elif yaqi_marks and player:
            # 选择最近的妖气标记
            yaqi_marks.sort(key=lambda x: abs(x[0] - player[0]) + abs(x[1] - player[1]))
            target = yaqi_marks[0]
            self.move_towards((target[0], target[1]))
            action_taken = True

        # 4. 如果没有目标且有传送门，则进入下一张图
        elif portal and not action_taken:
            self.move_towards(portal)
            time.sleep(1)
            pyautogui.press(self.keys["pick_item"])  # 进入传送门
            action_taken = True

        return frame

    def check_fatigue(self, frame):
        """检测疲劳值是否不足"""
        try:
            # 截取疲劳值区域（需要根据实际游戏界面调整坐标和大小）
            # 以下坐标为示例，请根据实际游戏界面进行调整
            fatigue_region = frame[10:30, 800:900]
            
            # 转换为HSV颜色空间以便于颜色识别
            hsv = cv2.cvtColor(fatigue_region, cv2.COLOR_BGR2HSV)
            
            # 定义疲劳值条的颜色范围（绿色）
            # 注意：这里的颜色范围需要根据实际游戏中的疲劳值条颜色进行调整
            lower_green = np.array([35, 50, 50])
            upper_green = np.array([85, 255, 255])
            
            # 创建掩码，只保留绿色区域
            mask = cv2.inRange(hsv, lower_green, upper_green)
            
            # 计算绿色区域的像素数量
            green_pixels = cv2.countNonZero(mask)
            
            # 计算疲劳值百分比
            total_pixels = fatigue_region.shape[0] * fatigue_region.shape[1]
            fatigue_percentage = (green_pixels / total_pixels) * 100
            
            print(f"当前疲劳值百分比: {fatigue_percentage:.2f}%")
            
            # 如果疲劳值低于阈值，返回True
            return fatigue_percentage < self.fatigue_threshold
        except Exception as e:
            print(f"检测疲劳值时出错: {str(e)}")
            return False

    def switch_character(self):
        """切换到下一个角色"""
        print("疲劳值不足，准备切换角色...")
        
        try:
            # 打开角色选择界面
            pyautogui.press(self.keys["character_select"])
            time.sleep(2)  # 等待角色选择界面加载
            
            # 选择下一个角色
            self.current_character_index = (self.current_character_index + 1) % len(self.characters)
            target_character = self.characters[self.current_character_index]
            print(f"切换到角色: {target_character} (索引: {self.current_character_index + 1}/{len(self.characters)})")
            
            # 根据角色索引选择角色（这里的坐标需要根据实际游戏界面调整）
            # 假设角色选择界面中，角色是横向排列的
            character_positions = [
                (400, 300),  # 第一个角色位置
                (600, 300),  # 第二个角色位置
                (800, 300),  # 第三个角色位置
                # 可以根据实际角色数量添加更多位置
            ]
            
            if self.current_character_index < len(character_positions):
                # 移动鼠标到目标角色位置并点击
                x, y = character_positions[self.current_character_index]
                pyautogui.moveTo(x, y)
                time.sleep(0.5)
                pyautogui.click()
                time.sleep(1)
                
                # 点击开始游戏按钮（坐标需要根据实际游戏界面调整）
                start_game_pos = (960, 540)  # 假设在屏幕中心
                pyautogui.moveTo(start_game_pos[0], start_game_pos[1])
                time.sleep(0.5)
                pyautogui.click()
                
                print(f"角色 {target_character} 切换成功，等待游戏加载...")
                time.sleep(10)  # 等待游戏加载完成
                return True
            else:
                print(f"未配置角色 {self.current_character_index} 的位置信息")
                return False
                
        except Exception as e:
            print(f"切换角色时出错: {str(e)}")
            return False

    def auto_farming_loop(self):
        """自动刷图主循环"""
        print("开始妖气追踪自动刷图...")
        print("按 'q' 键停止程序")

        while self.running:
            # 获取游戏画面
            frame = self.get_game_frame()
            if frame is None:
                time.sleep(0.1)
                continue

            # 检查疲劳值
            if self.check_fatigue(frame):
                # 疲劳值不足，切换角色
                if not self.switch_character():
                    print("所有角色疲劳值已耗尽，停止搬砖")
                    self.stop()
                    break
                continue

            # 检测妖气相关目标
            results = self.detect_yaqi_targets(frame)

            # 执行妖气追踪逻辑
            processed_frame = self.yaqi_tracking_logic(frame, results)

            # 显示处理后的画面
            cv2.imshow("DNF妖气追踪", cv2.resize(processed_frame, (800, 600)))

            # 检查退出指令
            if cv2.waitKey(1) & 0xFF == ord('q'):
                self.stop()
                break

            # 控制循环速度
            time.sleep(0.05)

    def start(self):
        """启动自动刷图"""
        self.running = True
        self.thread = threading.Thread(target=self.auto_farming_loop)
        self.thread.start()

    def stop(self):
        """停止自动刷图"""
        self.running = False
        if hasattr(self, 'thread') and self.thread.is_alive():
            self.thread.join()
        cv2.destroyAllWindows()
        print("妖气追踪自动刷图已停止")


def main():
    # 步骤1: 数据准备(首次运行时执行)
    dataset = DNFFiendishQiDataset()

    # 捕获游戏截图(手动操作游戏收集妖气场景数据)
    dataset.capture_in_game_screenshots(duration=500)  # 捕获10分钟

    # 标注数据后准备数据集(需要手动标注或使用标注工具)
    # dataset.prepare_dataset()

    # 创建配置文件
    # config_path = dataset.create_config_file()

    # 步骤2: 训练模型(数据准备完成后执行)
    # trainer = DNFFiendishQiTrainer(config_path)
    # trainer.train(epochs=120)  # 适当增加训练轮数提高精度

    # 步骤3: 运行自动妖气追踪刷图(模型训练完成后执行)
    # model_path = "dnf_yaqi_models/yaqi_tracker/weights/best.pt"
    # if os.path.exists(model_path):
    #     controller = DNFFiendishQiController(model_path)
    #     controller.start()
    #
    #     # 保持主程序运行
    #     try:
    #         while controller.running:
    #             time.sleep(1)
    #     except KeyboardInterrupt:
    #         controller.stop()
    # else:
    #     print(f"未找到模型文件: {model_path}")
    #     print("请先完成数据收集和模型训练步骤")


if __name__ == "__main__":
    main()
